# -*- coding:utf-8 -*-

from mitmproxy import ctx
import json
from pymongo import MongoClient
from datetime import datetime


class DB(object):
    def __init__(self):
        self.client = MongoClient('127.0.0.1', 27017)
        self.db = self.client["DouYin"]
        self.collection = self.db["search_keyword"]
    def insertItem(self, data):
        self.collection.insert_one(data)


db = DB()

# mitmdump.exe -p 8888 -s E:\Workspaces\Python\Spider\douyin\douyin_search_keyword_script.py
# 所有发出的请求数据包都会被这个方法所处理
# 所谓的处理，我们这里只是打印一下一些项；当然可以修改这些项的值直接给这些项赋值即可
def request(flow):
    # # 获取请求对象
    # request = flow.request
    # # 实例化输出类
    # info = ctx.log.info
    # # 打印请求的url
    # info(request.url)
    # # 打印请求方法
    # info(request.method)
    # # 打印host头
    # info(request.host)
    # # 打印请求端口
    # info(str(request.port))
    # # 打印所有请求头部
    # info(str(request.headers))
    # # 打印cookie头
    # info(str(request.cookies))
    pass

# 所有服务器响应的数据包都会被这个方法处理
# 所谓的处理，我们这里只是打印一下一些项
def response(flow):
    # 获取响应对象
    response = flow.response
    # # 实例化输出类
    # info = ctx.log.info
    # # 打印响应码
    # info(str(response.status_code))
    # # 打印所有头部
    # info(str(response.headers))
    # # 打印cookie头部
    # info(str(response.cookies))
    # # 打印响应报文内容
    # info(str(response.text))
    try:
        # 获取V用户信息GET
        # if 'aweme-eagle-hl.snssdk.com/aweme/v1/user/' in flow.request.url:
        if 'search-lf.amemv.com/aweme/v1/general/search/single' in flow.request.url:
            ctx.log.info("============获取数据=============")
            data_json = json.loads(response.text)
            for d in data_json['data']:
                if d.get('aweme_info') is not None:
                    tmp = d['aweme_info']['author']
                    dict_data = dict()
                    dict_data['enterprise_verify_reason'] = tmp.get('enterprise_verify_reason')  # 认证名称
                    dict_data['nickname'] = tmp.get('nickname')  # 昵称
                    dict_data['uid'] = tmp.get('uid')  # 用户ID
                    dict_data['unique_id'] = tmp.get('unique_id')  # 抖音号
                    dict_data['short_id'] = tmp.get('short_id')  # 抖音号

                    db.insertItem(dict_data)
                    # print(str(dict_data))
                    ctx.log.info("============Insert=============")
    except Exception as ex:
        ctx.log.warn("============错误=============")
        ctx.log.error(ex)